package stun2

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"goodlink/config"
	"goodlink/tools"
	"log"
	"math/big"
	"net"
	"time"

	"gogo"
)

func getStunIpPort5(attrType uint16, attributes []byte, attrLength uint16, magicCookie []byte, transactionID []byte) (string, int, error) {
	// https://www.rfc-editor.org/rfc/rfc5389.html#section-15.1
	// MAPPED-ADDRESS
	//  0                   1                   2                   3
	//   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |0 0 0 0 0 0 0 0|    Family     |           Port                |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |                                                               |
	// |                 Address (32 bits or 128 bits)                 |
	// |                                                               |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// https://www.rfc-editor.org/rfc/rfc5389.html#section-15.2
	// XOR-MAPPED-ADDRESS
	//  0                   1                   2                   3
	//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |x x x x x x x x|    Family     |         X-Port                |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |                X-Address (Variable)
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	attributeValue := attributes[4 : 4+attrLength]
	family := attributeValue[1]
	var ip, port []byte
	switch family {
	case 1:
		ip = attributeValue[4:8]
		port = attributeValue[2:4]
	case 2:
		ip = attributeValue[4:20]
	default:
		return "", 0, fmt.Errorf("unknown address family")
	}
	if attrType == 0x0020 { // XOR-Mapped Address
		for i := 0; i < 4; i++ {
			ip[i] ^= magicCookie[i]
		}
		if family == 2 {
			for i := 4; i < len(ip); i++ {
				ip[i] ^= transactionID[i-4]
			}
		}
	}
	port2 := binary.BigEndian.Uint16(port)
	return net.IP(ip).String(), int(port2), nil
}

func getStunIpPort4(response []byte, response_len int, needType uint16, magicCookie []byte, transactionID []byte) (string, int, error) {
	start := 0

	for {
		if start >= response_len-4 {
			break
		}

		// Parse STUN attributes
		attributes := response[start:]
		attrType := binary.BigEndian.Uint16(attributes[:2])
		attrLength := binary.BigEndian.Uint16(attributes[2:4])
		if attrLength < 8 {
			break
		}

		if attrType == needType {
			return getStunIpPort5(attrType, attributes, attrLength, magicCookie, transactionID)
		}

		start = start + int(attrLength) + 4
	}

	return "", 0, fmt.Errorf("attrType not found")
}

func getStunResponse(conn *net.UDPConn, addr string, buf *bytes.Buffer) ([]byte, int, error) {
	udpAddr, err := net.ResolveUDPAddr("udp4", addr)
	if err != nil {
		return nil, 0, err
	}

	_, err = conn.WriteToUDP(buf.Bytes(), udpAddr)
	if err != nil {
		return nil, 0, err
	}

	response := make([]byte, 1024)

	conn.SetReadDeadline(time.Now().Add(1000 * time.Millisecond))
	n, err := conn.Read(response)
	defer conn.SetDeadline(time.Time{})

	if err != nil {
		return nil, 0, err
	}
	if n < 32 {
		return nil, 0, fmt.Errorf("invalid response")
	}

	// Parse STUN message
	if !bytes.Equal(response[4:8], buf.Bytes()[4:8]) {
		return nil, 0, fmt.Errorf("invalid magic cookie in response")
	}
	if !bytes.Equal(response[8:20], buf.Bytes()[8:20]) {
		return nil, 0, fmt.Errorf("transaction ID mismatch in response")
	}

	return response, n, nil
}

func getStunIpPort2(conn *net.UDPConn, addr string) (string, int, int, error) {

	// https://www.rfc-editor.org/rfc/rfc5389.html#section-6
	// STUN Message Structure
	//   0                   1                   2                   3
	//   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |0 0|     STUN Message Type     |         Message Length        |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |                         Magic Cookie                          |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |                                                               |
	// |                     Transaction ID (96 bits)                  |
	// |                                                               |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

	// STUN message header
	buf := new(bytes.Buffer)
	// Start with fixed 0x00, message type: 0x01, message length: 0x0000
	buf.Write([]byte{0x00, 0x01, 0x00, 0x00})
	magicCookie := []byte{0x21, 0x12, 0xA4, 0x42}
	buf.Write(magicCookie)
	transactionID := make([]byte, 12)
	rand.Read(transactionID)
	buf.Write(transactionID)

	response, n, err := getStunResponse(conn, addr, buf)
	if err != nil {
		return "", 0, 0, err
	}

	// https://www.rfc-editor.org/rfc/rfc5389.html#section-15
	// STUN Attributes
	//   0                   1                   2                   3
	//   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |         Type                  |            Length             |
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	// |                         Value (variable)                ....
	// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

	wan_ip1, wan_port1, err := getStunIpPort4(response[20:], n, 0x0001, magicCookie, transactionID)
	if err != nil {
		wan_ip1, wan_port1, err = getStunIpPort4(response[20:], n, 0x0020, magicCookie, transactionID)
	}
	if err != nil {
		return "", 0, 0, err
	}

	var change_ip, wan_ip2 string
	var change_port, wan_port2 int

	change_ip, change_port, err = getStunIpPort4(response[20:], n, 0x0005, magicCookie, transactionID)
	if err == nil {
		response, n, err = getStunResponse(conn, fmt.Sprintf("%s:%d", change_ip, change_port), buf)
		if err == nil {
			wan_ip2, wan_port2, err = getStunIpPort4(response[20:], n, 0x0001, magicCookie, transactionID)
			if err != nil {
				wan_ip2, wan_port2, err = getStunIpPort4(response[20:], n, 0x0020, magicCookie, transactionID)
			}
			if err == nil {
				if wan_ip1 != wan_ip2 {
					gogo.Log().ErrorF("wan_ip: %s, %s", wan_ip1, wan_ip2)
				}
				return wan_ip1, wan_port1, wan_port2, nil
			}
		}
		return wan_ip1, wan_port1, 0, err
	}

	return wan_ip1, wan_port1, 0, nil
}

func GetWanIpPort2(conn *net.UDPConn) (string, int, int) {
	gogo.Log().Debug("   获取本端地址")
	defer conn.SetReadDeadline(time.Time{})
	n2 := 0

	for {
		n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(config.GetConfig().StunList))))
		stun_svr := config.GetConfig().StunList[n.Int64()]
		wan_ip, wan_port1, wan_port2, err := getStunIpPort2(conn, stun_svr)
		//log.Printf("   stun_svr: %s, wan_ip: %s, wan_port1: %d, wan_port2: %d, err: %v", stun_svr, wan_ip, wan_port1, wan_port2, err)
		if err != nil {
			log.Printf("   stun_svr: %s, err: %v", stun_svr, err)
			/*if strings.Contains(err.Error(), "timeout") || strings.Contains(err.Error(), "closed") {
				time.Sleep(1 * time.Second)
			}*/
		}
		if wan_ip != "" && wan_port1 > 0 {
			return wan_ip, wan_port1, wan_port2
		}
		n2++
		if n2 >= len(config.GetConfig().StunList) {
			time.Sleep(3 * time.Second)
			n2 = 0
		}
	}
}

func GetWanIpPort() (string, int, int) {
	conn := tools.GetListenUDP()
	defer conn.Close()
	return GetWanIpPort2(conn)
}

func TestStun() {
	conn := tools.GetListenUDP()

	for {
		for _, stun_svr := range config.GetConfig().StunList {
			conn.SetReadDeadline(time.Now().Add(1000 * time.Millisecond))
			if wan_ip, wan_port1, wan_port2, _ := getStunIpPort2(conn, stun_svr); wan_ip != "" && wan_port1 > 0 && wan_port2 > 0 {
				conn.SetReadDeadline(time.Time{})
				fmt.Printf("stun_svr: %s, wan_ip: %s, wan_port1: %d, wan_port2: %d\n", stun_svr, wan_ip, wan_port1, wan_port2)
			} else {
				fmt.Printf("stun_svr: %s, failed\n", stun_svr)
			}
		}
		time.Sleep(3 * time.Second)
	}
}
